package org.javaagent;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;

import java.io.ByteArrayInputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

/**
 * @author shenguangyang
 */
public class CostTransformer implements ClassFileTransformer {
    private final String packageName;

    public CostTransformer(String packageName) {
        this.packageName = packageName;
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        className = className.replace("/", ".");
        // 只对目标包下进行耗时统计, $$ 表示当前类是代理对象
        if (!className.startsWith(packageName) || className.contains("$$")) {
            return classfileBuffer;
        }

        CtClass cl;
        try {
            ClassPool classPool = ClassPool.getDefault();
            cl = classPool.makeClass(new ByteArrayInputStream(classfileBuffer));
            System.out.println("className: " + className);
            for (CtMethod method : cl.getDeclaredMethods()) {
                // 所有方法, 统计耗时
                // 注意: 需要通过addLocalVariable来声明局部变量
                method.addLocalVariable("start", CtClass.longType);
                method.insertBefore("start = System.currentTimeMillis();");
                String methodName = method.getLongName();
                // 请注意第二个参数，设置为true，则表示即便抛出异常了，下面的代码也会执行；相当于将它封装在finally里面了
                method.insertAfter(
                        "System.out.println(\"" + methodName + " cost: \" + (System.currentTimeMillis() - start) + \" ms\");");
            }
            return cl.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return classfileBuffer;
    }
}
